


import logging
import cv2
import time
import asyncio
import op_route
import keti4_config as config

class Dispatcher:
    def __init__(self, send_port=8800, target_nodes=[0x02, 0x03, 0x04, 0x05], file="data.png"):
        self.file = file
        self.send_port = send_port
        self.target_nodes = target_nodes

    # 主调度逻辑：读取图像、分割并发送
    async def run(self):
        logging.info("=== Dispatcher ===")
        logging.info(f"<run>Dispatcher info: send_port={self.send_port}, target_nodes={self.target_nodes}, file={self.file}")

        # 1. 读取图像
        logging.info("<run>start read image")
        image = self.read_image()

        # 2. 图像分割
        logging.info("<run>start split image")
        image_list = self.split_image(image)

        # 3. 并发发送图像到辅星
        tasks = []
        logging.info("<run>create async tasks for sending images to sub stars")
        for idx in range(len(image_list)):
            tasks.append(self.send_image_to_sub_star(image_list[idx], 
                                                     target_node=self.target_nodes[idx], 
                                                     send_port=self.send_port))

        # 异步执行所有任务
        await asyncio.gather(*tasks)

        logging.info("__________Dispatcher Finish______________") 
        return True

    def read_image(self) -> bytes:
        logging.info(f"[Dispatcher]<read_image> file_name: {self.file}")

        """读取图像"""
        logging.info(f"[Dispatcher]<read_image> file_name: {self.file}")
        image = cv2.imread(self.file)
        if image is None:
            logging.error(f"[Dispatcher]<read_image> 读取图像失败: {self.file}")
            raise FileNotFoundError(f"无法读取图像: {self.file}")
        logging.info(f"[Dispatcher]<read_image> 图像读取成功, shape={image.shape}")
        return image

        # return b'123456789'  # 模拟图像数据

#     def split_image(self, image, n=4) -> list:
#         logging.info(f"[Dispatcher]<split image>: len {len(image)}")
#         # 模拟图像分割
#         # return [b'part1', b'part2', b'part3', b'part4']


# --- 分割图像 ---
#  # 验证分割结果
#     for idx, block in enumerate(split_blocks):
#         cv2.imwrite(f"block_{idx}.png", block)  # 将每个分块保存成图片


    def split_image(self, image, n=4) -> list:
        """
        将图像均匀分割成 n x n 块

        :param image: 输入的图像（numpy.ndarray）
        :param n: 分割数（每行/每列分割成 n 块）
        :return: 分割后的图像块列表
        """
        if n <= 0:
            logging.error("[Dispatcher]<split_image> 分割数 n 必须大于 0")
            raise ValueError("分割数 n 必须大于 0")

        height, width, channels = image.shape
        logging.info(f"[Dispatcher]<split_image> 开始分割图像，分割数: {n}, 原图尺寸: {width}x{height}")

        # 计算每个小块的尺寸
        block_h = height // n
        block_w = width // n

        image_blocks = []

        # 遍历分割图像
        for i in range(n):
            for j in range(n):
                # 计算当前块的位置
                y_start = i * block_h
                y_end = (i + 1) * block_h if i != n - 1 else height  # 确保最后一块不会丢失像素
                x_start = j * block_w
                x_end = (j + 1) * block_w if j != n - 1 else width

                # 截取图像块
                block = image[y_start:y_end, x_start:x_end]
                image_blocks.append(block)
                logging.info(f"[Dispatcher]<split_image> 分割块[{i},{j}] 尺寸: {block.shape}")

        logging.info(f"[Dispatcher]<split_image> 图像分割完成，总共 {len(image_blocks)} 块")
        return image_blocks

    # 异步发送图像到辅星
    async def send_image_to_sub_star(self, image, target_node=0x02, send_port=8800):
        data = bytes(image)
        logging.info(f"[Dispatcher]<send_image_to_sub_star> datalen: {len(data)}, target_node: {target_node}, send_port: {send_port}")

        # for i in range(10):
        #     await asyncio.sleep(1)  # 异步等待，不阻塞事件循环
        #     logging.info(f"[Dispatcher]<send_image_to_sub_star> sending to target_node: {target_node} ... {i}")

        # 初始化路由器客户端
        router = op_route.RouterClient(host='localhost', port=send_port)
        router.connect()
        data_bytes = bytes(data)

        # 发送数据
        router.send_packet(op_route.DATA_TYPE_SEND, target_node, len(data_bytes), data_bytes)
        logging.info("Sent data:", data)



if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # 创建 Dispatcher 实例
    dispatcher = Dispatcher(send_port=config.send_port, target_nodes=[0x02, 0x03, 0x04, 0x05], file="data.png")

    # 运行 Dispatcher
    asyncio.run(dispatcher.run())
